from torch import nn

from dataset import *
from positional_encoding import PositionalEncoding
from mask import generate_square_subsequent_mask


class Seq2seqTransformer(nn.Module):
    def __init__(self,
                 src_vocab_size: int,
                 tgt_vocab_size: int,
                 d_model: int = 512,
                 n_heads: int = 8,
                 num_encoder_layers: int = 6,
                 num_decoder_layers: int = 6,
                 dim_feedforward: int = 2048,
                 dropout: float = 0.1,
                 max_len: int = 5000):
        """
        src_vocab_size: 源词汇表大小
        tgt_vocab_size: 目标词汇表大小
        d_model: 模型特征维度
        n_heads: 注意力头数
        num_encoder_layers: 编码器层数
        num_decoder_layers: 解码器层数
        dim_feedforward: 前馈网络维度
        dropout: Dropout比率
        max_len: 序列最大长度
        """
        # 调用父类 nn.Module 的初始化方法
        super(Seq2seqTransformer, self).__init__()
        # 创建源序列词嵌入层，将源词汇索引映射为向量表示
        self.src_embedding = nn.Embedding(src_vocab_size, d_model, padding_idx=PAD_IDX)
        # 创建目标序列词嵌入层，将目标词汇索引映射为向量表示
        self.tgt_embedding = nn.Embedding(tgt_vocab_size, d_model, padding_idx=PAD_IDX)
        # 创建位置编码模块，为序列添加位置信息
        self.pos_encoding = PositionalEncoding(d_model, max_len, dropout)
        # 创建PyTorch原生Transformer模块，包含编码器和解码器
        self.trans = nn.Transformer(
            d_model=d_model,  # 模型特征维度
            nhead=n_heads,  # 注意力头数
            num_encoder_layers=num_encoder_layers,  # 编码器层数
            num_decoder_layers=num_decoder_layers,  # 解码器层数
            dim_feedforward=dim_feedforward,  # 前向传播隐藏层维度
            dropout=dropout,  # Dropout比率
            batch_first=True  # Batch_Size顺序
        )
        # 创建输出层，将模型输出映射到目标词汇表空间
        self.generator = nn.Linear(d_model, tgt_vocab_size)

    """
    掩码的基本作用
    掩码在Transformer模型中主要用于：
        屏蔽无效信息：隐藏填充符号(PAD_IDX)对模型的影响
        防止信息泄露：在解码时防止看到未来位置的信息
        提高模型效率：避免对无意义位置的计算
    src_key_padding_mask 和 tgt_key_padding_mask
        作用：标记序列中的填充值位置
        值：True表示填充值，False表示有效值
    tgt_mask (因果掩码)
        作用：防止解码器在训练时看到未来位置的信息
        值：上三角矩阵，对角线以上为-inf，其余为0
    掩码在注意力机制中的应用
        在注意力计算中，掩码会加到注意力分数上：
        这种设计确保了模型在处理变长序列时能够正确关注有效信息，同时在解码过程中保持自回归特性。
    """

    def forward(self,  # 源序列、目标序列、源序列填充掩码、目标序列填充掩码、目标序列掩码
                src: torch.Tensor,  # 源序列
                tgt_input: torch.Tensor,  # 目标序列输入
                src_key_padding_mask: torch.Tensor,  # 源序列填充掩码
                tgt_key_padding_mask: torch.Tensor,  # 目标序列填充掩码
                tgt_mask: torch.Tensor):  # 目标序列因果掩码
        """
        定义前向传播方法，接收源序列、目标序列及相关掩码
        """
        # 对源序列进行词嵌入，然后添加位置编码
        src_embedding = self.pos_encoding(self.src_embedding(src))
        # 对目标序列进行词嵌入，然后添加位置编码
        tgt_embedding = self.pos_encoding(self.tgt_embedding(tgt_input))
        # 将处理后的序列输入Transformer模块进行计算
        out = self.trans(
            src=src_embedding,  # 传入源序列的嵌入表示（已包含词嵌入和位置编码），形状为 [batch_size, src_seq_len, d_model]
            tgt=tgt_embedding,  # 传入目标序列的嵌入表示（已包含词嵌入和位置编码），形状为 [batch_size, tgt_seq_len, d_model]
            src_key_padding_mask=src_key_padding_mask,  # 传入源序列的填充掩码，标记源序列中的填充值位置，形状为 [batch_size, src_seq_len]
            tgt_key_padding_mask=tgt_key_padding_mask,  # 入目标序列的填充掩码，标记目标序列中的填充值位置，形状为 [batch_size, tgt_seq_len]
            memory_key_padding_mask=src_key_padding_mask,  # 传入记忆（编码器输出）的填充掩码，用于解码器注意力计算，通常与源序列填充掩码相同
            tgt_mask=tgt_mask  # 传入目标序列的因果掩码（上三角矩阵），防止解码器在训练时看到未来位置的信息
        )
        # 通过线性层生成最终的词汇分布输出
        return self.generator(out)

    """
    Seq2seqTransformer 类中的 greedy_decode 方法，用于执行贪心解码（greedy decoding）来生成序列。
    主要功能
        序列生成：根据输入的源序列 src 生成目标序列
        贪心策略：每一步都选择概率最高的词汇作为下一个输出
        自回归解码：逐步生成输出序列，每一步的生成都依赖于之前生成的词汇
    工作流程
    1.初始化阶段：
        设置模型为评估模式
        对源序列进行嵌入和位置编码
        通过编码器生成记忆向量 memory
        初始化输出序列 ys 为开始标记 BOS_IDX
    2.迭代解码：
        在每次迭代中，对当前输出序列进行嵌入和位置编码
        生成因果掩码 tgt_mask 防止看到未来信息
        通过解码器生成输出
        使用 generator 将输出映射到词汇表空间
        选择概率最高的词汇作为下一个输出
    3.终止条件：
        达到最大长度 max_len
        所有样本都生成了结束标记 EOS_IDX
    关键特性
        使用 @torch.no_grad() 装饰器禁用梯度计算，提高推理效率
        支持批量处理多个序列
        处理变长序列的填充掩码
        实现提前终止机制，当所有序列都完成时提前退出
    """

    @torch.no_grad()
    def greedy_decode(self, src: torch.Tensor, max_len: int = 32) -> torch.Tensor:
        """
        定义贪心解码方法，接收源序列和最大长度，返回生成的目标序列
        src:[B,S]
        """
        # 将模型设置为评估模式，关闭如 dropout 等训练时特有的行为
        self.eval()
        # 获取批次大小 B
        B = src.size(0)

        # 创建源序列的填充掩码，标记哪些位置是填充值 PAD_IDX
        src_pad = (src == PAD_IDX)  # 等于PAD_IDX的元素为True, 否则为False
        # 对源序列进行词嵌入 (self.src_embedding)，然后添加位置编码 (self.pos_encoding)
        src_embedding = self.pos_encoding(self.src_embedding(src))  # [B,S,D]
        # 通过 self.trans.encoder 处理源序列，生成编码记忆 memory 供解码器使用
        memory = self.trans.encoder(src_embedding, src_key_padding_mask=src_pad)

        # 初始化输出序列，每个样本以开始标记 BOS_IDX 开始
        ys = torch.full((B, 1), BOS_IDX, dtype=torch.long, device=src.device)  # [B,1]
        # 创建完成状态张量，记录每个样本是否已完成生成（遇到结束标记）
        finished = torch.zeros(B, dtype=torch.bool, device=src.device)

        # 最多进行 max_len - 1 次迭代（减1是因为已有一个 BOS_IDX）
        for _ in range(max_len - 1):
            # 获取当前输出序列的长度 T
            T = ys.size(1)
            # 对当前输出序列进行词嵌入和位置编码,
            tgt = self.pos_encoding(self.tgt_embedding(ys))  # [B,T,D]

            # 生成因果掩码，防止解码器看到未来位置的信息
            tgt_mask = generate_square_subsequent_mask(T).to(src.device)  # [T,T], -inf 上三角
            # 创建目标序列的填充掩码
            tgt_pad = (ys == PAD_IDX)

            # 通过 self.trans.decoder 生成输出，传入目标序列、编码记忆和各种掩码
            out = self.trans.decoder(
                tgt=tgt,  # [B,T,D] 目标序列的嵌入表示
                memory=memory,  # [B,S,D] 编码器的输出
                tgt_mask=tgt_mask,  # [T,T] 目标序列的因果掩码
                tgt_key_padding_mask=tgt_pad,  # [B,T] 目标序列的填充掩码
                memory_key_padding_mask=src_pad  # [B,S] 编码器的填充掩码
            )  # [B,T,D]

            # 取最后一个时间步的输出，通过生成器 (self.generator) 映射到词汇表空间
            logits = self.generator(out[:, -1, :])  # [B,V]
            # 选择概率最高的词汇作为下一个生成的 token
            next_tok = logits.argmax(dim=-1, keepdim=True)  # [B,1]

            # 对已完成的序列，强制其继续生成结束标记 EOS_IDX
            next_tok = torch.where(
                finished.unsqueeze(1),
                torch.full_like(next_tok, EOS_IDX),
                next_tok
            )
            # 将新生成的 token 添加到输出序列中
            ys = torch.cat([ys, next_tok], dim=1)  # [B,T+1]
            # 更新完成状态，标记新完成的序列
            finished |= (next_tok.squeeze(1) == EOS_IDX)
            # 如果所有序列都已完成，则提前结束循环
            if torch.all(finished):
                break
        # 返回完整的生成序列
        return ys
